﻿using System;
using System.Collections.Generic;
using System.Linq;

namespace Algoritmia.DataStructures.PriorityQueues
{

    public class Heap<T> : 
        IPriorityQueue<T>, 
        IEnumerable<T>
    {
        private static readonly IEnumerable<T> DefaultElement = new[] {default(T)};
        protected readonly Func<T,T,bool> _firstHasHigherPriority;
        protected T[] _heap;
        protected int _count;
        protected int _capacity;

        public Heap(IEnumerable<T> data = null, int capacity = 0, Func<T, T, bool> firstHasHigherPriority = null, Comparison<T> comparison =null)
        {
            _firstHasHigherPriority = GetPriorityFunction(comparison, firstHasHigherPriority);
            _heap = DefaultElement.Concat(data ?? Enumerable.Empty<T>()).ToArray();
            _count = _heap.Length - 1;
            _capacity = Math.Max(_count, _capacity == 0 ? 1 : capacity);
            if (_capacity > _count)
            {
                Array.Resize(ref _heap, _capacity + 1);
            }
            BuildHeap();
        }

        private void BuildHeap()
        {
            for(var i = _count/2; i > 0; i--)
            {
                Heapify(i);
            }
        }

        private static Func<T, T, bool> GetPriorityFunction(Comparison<T> comparer, Func<T, T, bool> firstHasHigherPriority)
        {
            if (comparer != null && firstHasHigherPriority != null)
            {
                throw new ArgumentException("You can't define firstHasHigherPriority and comparer simultaneously");
            }
            if (comparer != null)
            {
                return (first, second) => comparer(first, second) < 0;
            }
            if (firstHasHigherPriority == null)
            {
                if (typeof(T).GetInterfaces().Contains(typeof(IComparable<T>)))
                {
                    return (first, second) => (Comparer<T>.Default.Compare(first,second) > 0);
                }
                throw new ArgumentException("No comparison criterium on non comparable data");
            }
            return firstHasHigherPriority;
        }

        protected void Swap(int i, int j)
        {
            var aux = _heap[i];
            _heap[i] = _heap[j];
            _heap[j] = aux;
        }

        protected virtual void Heapify(int index)
        {
            var best = index;
            while (true)
            {
                var left = 2 * index;
                var right = 2 * index + 1;
                if (left <= _count && _firstHasHigherPriority(_heap[left], _heap[index]))
                {
                    best = left;
                }
                if (right <= _count && _firstHasHigherPriority(_heap[right],_heap[best]))
                {
                    best = right;
                }
                if (best == index) break;

                Swap(index, best);
                index = best;
            }
        }

        protected virtual void BubbleUp(int index)
        {
            int parent = index / 2;
            while (index > 1 && _firstHasHigherPriority(_heap[index], _heap[parent]))
            {
                Swap(parent, index);
                index = parent;
                parent /= 2;
            }
        }

        private void DoubleHeapCapacity()
        {
            _capacity *= 2;
            Array.Resize(ref _heap, _capacity + 1);
        }


        #region IPriorityQueue<T> members

        public virtual void Add(T elem)
        {
            _count++;
            if (_count >= _capacity)
            {
                DoubleHeapCapacity();
            }
            _heap[_count] = elem;
            BubbleUp(_count);
        }

        public virtual T ExtractOptimal()
        {
            if (_count == 0)
                throw new InvalidOperationException("Can't extract anything from an empty heap.");
            T value = _heap[1];
            _heap[1] = _heap[_count];
            _count--;
            if ( _count > 1 ) Heapify(1);
            return value;
        }

        public virtual T GetOptimal()
        {
            if (_count == 0)
                throw new InvalidOperationException("Can't get anything from an empty heap.");
            return _heap[1];
        }

        #endregion

        #region ICountable members

        public int Count
        {
            get { return _count; }
        }

        #endregion

        #region IEnumerable<T> members

        public IEnumerator<T> GetEnumerator()
        {
            int i = 1;
            while (i <= _count)
                yield return _heap[i++];
        }

       
        System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator()
        {
            return GetEnumerator();
        }

        #endregion

        public override string ToString()
        {
            return String.Format("{0} <{1}> [{2}] optimalFunction = {3}", GetType().Name, typeof(T).Name,
                                    string.Join(",", this), _firstHasHigherPriority);
        }

    }
}
